/**
 * Copyright (C), 2020-2021 by colley.ma
 * File Name: RedisLimitRater.java
 * Encoding: UTF-8
 * Date: 2021年3月26日 下午3:48:28
 * History:
*/
package com.swak.frame.ratelimit.resolver;

import java.time.Instant;
import java.util.Arrays;
import java.util.List;

import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.cache.annotation.CachingConfigurationSelector;
import org.springframework.context.ApplicationListener;
import org.springframework.core.io.ClassPathResource;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.util.ClassUtils;

import com.swak.frame.cache.CacheProxy;
import com.swak.frame.cache.OutputType;
import com.swak.frame.common.SpringBeanFactory;
import com.swak.frame.config.AbstractStateConfigurable;
import com.swak.frame.exception.SwakAssert;
import com.swak.frame.exception.SwakScriptNotFoundException;
import com.swak.frame.ratelimit.LimitType;
import com.swak.frame.ratelimit.SwakRateLimiter;
import com.swak.frame.ratelimit.config.LimitCommand;
import com.swak.frame.ratelimit.event.RateAlarmEvent;
import com.swak.frame.ratelimit.event.RateLimitEvent;
import com.swak.frame.ratelimit.event.RateLimitEventInfo;
import com.swak.frame.script.DefaultRedisScript;
import com.swak.frame.script.RedisScript;

public class RedisLimitRater extends AbstractStateConfigurable<LimitCommand>
		implements SwakRateLimiter, ApplicationListener<RateLimitEvent>, InitializingBean {

	private Logger logger = LoggerFactory.getLogger(getClass()); // 日志记录

	private CacheProxy cacheProxy;

	private String luaString;

	private RedisScript script;

	public RedisLimitRater(CacheProxy cacheProxy) {
		super(LimitCommand.class);
		this.cacheProxy = cacheProxy;
	}

	@Override
	public boolean tryAcquire(String resource) {
		if (StringUtils.isEmpty(resource)) {
			return true;
		}
		LimitCommand config = getConfig().getOrDefault(resource, new LimitCommand(false));
		if (!config.isEnable() || config.getQps() <= 0) {
			return true;
		}
		return isAllowed(config, resource);
	}

	@Override
	public void afterPropertiesSet() throws Exception {
		SwakAssert.notNull(cacheProxy,"cacheProxy is null!");
		DefaultRedisScript redisScript = new DefaultRedisScript();
		redisScript.setScriptSource(
				new ResourceScriptSource(new ClassPathResource("META-INF/scripts/redis/rate_limiter.lua")));
		this.script = redisScript;
		this.luaString = cacheProxy.scriptLoad(this.script.getSha1());
	}

	static List<String> getKeys(String resource) {
		// use `{}` around keys to use Redis Key hash tags
		// this allows for using redis cluster
		// Make a unique key per user.
		String prefix = "swak_rate_limiter.{" + resource;

		// You need two Redis keys for Token Bucket.
		String tokenKey = prefix + "}.tokens";
		String timestampKey = prefix + "}.timestamp";
		return Arrays.asList(tokenKey, timestampKey);
	}

	public boolean isAllowed(LimitCommand config, String resource) {
		List<String> scriptArgs = Arrays.asList(config.getQps() + "", config.getCapacity() + "",
				Instant.now().getEpochSecond() + "", "1");
		Object result = null;
		try {
			try {
				result = cacheProxy.evalsha(luaString, getKeys(resource), scriptArgs, false, OutputType.MULTI);
			} catch (SwakScriptNotFoundException e) {
				this.luaString = cacheProxy.scriptLoad(script.getSha1());
				result = cacheProxy.evalsha(luaString, getKeys(resource), scriptArgs, false, OutputType.MULTI);
			}
		} catch (Exception e) {
			logger.error("Error determining if user allowed from redis,resource:" + resource, e);
		}
		if (result != null) {
			@SuppressWarnings("unchecked")
			List<Long> response = (List<Long>) result;
			if (CollectionUtils.isNotEmpty(response)) {
				boolean acquire = response.get(0) == 1L;
				if (!acquire) {
					if (config != null && config.getEvent()!=null) {
						RateAlarmEvent alarmEvent = SpringBeanFactory.getBean(config.getEvent());
						if (alarmEvent != null) {
							RateLimitEventInfo eventInfo = RateLimitEventInfo.newRateLimitEventInfo(resource, config,LimitType.REDIS);
							alarmEvent.event(eventInfo);
						}
					}
				}
				return acquire;
			}
		}
		return true;
	}

	@Override
	public void onApplicationEvent(RateLimitEvent event) {
		LimitCommand limitConfig = event.getConfig();
		if (LimitType.REDIS.equals(event.getLimitType())) {
			if (limitConfig != null) {
				limitConfig.setDynamic(false);//分布是缓存不走动态分配
				getConfig().put(event.getResource(), limitConfig);
			}
		}
	}
	
	public static void main(String[] args) {
		ClassLoader classLoader = CachingConfigurationSelector.class.getClassLoader();
		boolean jsr107Present = ClassUtils.isPresent("javax.cache.Cache", classLoader);
		System.out.println(jsr107Present);
	}
}
